#!/usr/bin/env python
#Adapted from https://gist.github.com/batzner/7c24802dd9c5e15870b4b56e22135c96

import tensorflow as tf

def rename(checkpoint_dir, replace_from, replace_to, add_prefix, dry_run=False):
    checkpoint = tf.train.get_checkpoint_state(checkpoint_dir)
    with tf.Session() as sess:
        for var_name, _ in tf.contrib.framework.list_variables(checkpoint_dir):
            # Load the variable
            var = tf.contrib.framework.load_variable(checkpoint_dir, var_name)

            # Set the new name
            new_name = var_name
            if None not in [replace_from, replace_to]:
                new_name = new_name.replace(replace_from, replace_to)
            if add_prefix:
                new_name = add_prefix + new_name

            if dry_run:
                print(('%s would be renamed to %s.' % (var_name, new_name)))
            else:
                print(('Renaming %s to %s.' % (var_name, new_name)))
                # Rename the variable
                var = tf.Variable(var, name=new_name)

        if not dry_run:
            # Save the variables
            saver = tf.train.Saver()
            sess.run(tf.global_variables_initializer())
            # saver.save(sess, checkpoint.model_checkpoint_path)
            saver.save(sess, checkpoint_dir)

#checkpoint_dir="/home/dlrc/dlrc17-gdk/ext/OnAVOS/models/small_lego2-00012453"
#checkpoint_dir="models/COCO_SPN_exp2_with_prefix/COCO_SPN_exp2-00000151"
checkpoint_dir="models/SPN_exp_with_prefix/SPN_exp-00000200"
replace_from=None
replace_to=None
add_prefix="SPN/"

rename(checkpoint_dir,replace_from,replace_to,add_prefix)
